import json
import os
import cv2
import shutil

import numpy as np

frames_dirs = [
    'C:/Users/jario/Downloads/driver_161_90frame/06030819_0755.MP4',
    'C:/Users/jario/Downloads/driver_161_90frame/06030831_0759.MP4',
]
gray_dir = 'C:/dataset/A2RL/ws_lane_240403/Yas_Grey_240322'


def show_lane_in_dir(frames_dir):
    file_names = os.listdir(frames_dir)
    father_dir = os.path.dirname(frames_dir)
    base_name = os.path.basename(frames_dir)
    line_gap = 10

    json_file = os.path.join(father_dir, 'label_data_' + base_name + '.json')
    json_ff = open(json_file, 'w')

    colors = [(0, 0, 255), (0, 255, 255), (255, 0, 0), (0, 255, 0)]

    for file_name in file_names:
        if os.path.splitext(file_name)[-1] == '.jpg':
            img_fn = os.path.join(frames_dir, file_name)
            ann_fn = os.path.join(frames_dir, os.path.splitext(file_name)[0] + '.lines.txt')
            img = cv2.imread(img_fn)
            # print(img_fn)
            img2 = img.copy()
            lanes = []
            if not os.path.exists(ann_fn):
                os.remove(img_fn)
                continue
            with open(ann_fn) as f:
                lines = f.readlines()
                for line in lines:
                    line = [int(round(float(l))) for l in line.strip().split(' ') if len(l) > 0]
                    lane = []
                    for i in range(0, len(line), 2):
                        lane.append([line[i], line[i + 1]])
                    lanes.append(lane)

            need_size = [1280, 720]
            scale = need_size[0] / img.shape[1]
            img = cv2.resize(img, (need_size[0], int(need_size[0] / img.shape[1] * img.shape[0])))

            if img.shape[0] < need_size[1]:
                margin = int((need_size[1] - img.shape[0]) / 2)
                img3 = np.ones((need_size[1], need_size[0], 3), dtype=np.uint8) * 127
                img3[margin:-margin, :] = img
                img = img3
                new_lanes = []
                for i, lane in enumerate(lanes):
                    new_lane2 = []
                    for pt in lane:
                        new_pt2 = [int(pt[0] * scale), int(pt[1] * scale + margin)]
                        new_lane2.append(new_pt2)
                    new_lanes.append(new_lane2)
                lanes = new_lanes
            img4 = img.copy()

            binary_image1 = np.zeros((img.shape[0], img.shape[1], 1), np.uint8)
            binary_image2 = binary_image1.copy()
            binary_image3 = binary_image1.copy()
            binary_image4 = binary_image1.copy()
            binary_image_v2 = np.zeros((img.shape[0], img.shape[1], 1), np.uint8)
            for i, lane in enumerate(lanes):
                for pt in lane:
                    cv2.circle(img, pt, 3, colors[i], -1, cv2.LINE_AA)
                    cv2.circle(img, pt, 3, (0, 0, 0), 1, cv2.LINE_AA)
                for j in range(len(lane)):
                    """
                    min_dist = 1e8
                    pt_k = None
                    for k in range(len(lane)):
                        if j != k:
                            dist = np.sqrt((lane[j][0]-lane[k][0]) ** 2 + (lane[j][1]-lane[k][1]) ** 2)
                            if dist < min_dist:
                                min_dist = dist
                                pt_k = lane[k]
                    """
                    if j > 0:
                        cv2.line(img, lane[j-1], lane[j], (255, 255, 255), 1)
                        if i == 0:
                            cv2.line(binary_image1, (lane[j - 1][0] - 8, lane[j - 1][1]), (lane[j][0] - 8, lane[j][1]),
                                     (50,), 1)
                            cv2.line(binary_image2, (lane[j - 1][0] + 8, lane[j - 1][1]), (lane[j][0] + 8, lane[j][1]),
                                     (100,), 1)
                            cv2.line(binary_image_v2, (lane[j - 1][0] - 8, lane[j - 1][1]), (lane[j][0] - 8, lane[j][1]),
                                     (1,), 2)
                            cv2.line(binary_image_v2, (lane[j - 1][0] + 8, lane[j - 1][1]), (lane[j][0] + 8, lane[j][1]),
                                     (2,), 2)
                        else:
                            cv2.line(binary_image4, (lane[j - 1][0] + 8, lane[j - 1][1]), (lane[j][0] + 8, lane[j][1]),
                                     (200,), 1)
                            cv2.line(binary_image3, (lane[j - 1][0] - 8, lane[j - 1][1]), (lane[j][0] - 8, lane[j][1]),
                                     (150,), 1)
                            cv2.line(binary_image_v2, (lane[j - 1][0] + 8, lane[j - 1][1]), (lane[j][0] + 8, lane[j][1]),
                                     (4,), 2)
                            cv2.line(binary_image_v2, (lane[j - 1][0] - 8, lane[j - 1][1]), (lane[j][0] - 8, lane[j][1]),
                                     (3,), 2)

            lanes = []
            h_samples = []
            lane1 = []
            lane2 = []
            lane3 = []
            lane4 = []
            for i in range(160, binary_image1.shape[0], line_gap):
                has_lane = False
                for j in range(binary_image1.shape[1]):
                    if binary_image1[i, j] == 50 or binary_image2[i, j] == 100 or binary_image3[i, j] == 150 or binary_image4[i, j] == 200:
                        has_lane = True
                if not has_lane:
                    continue
                # if not has_lane1:
                #     assert False, 'why not has lane1 {}'.format(file_name)
                # if not has_lane2:
                #     assert False, 'why not has lane2 {}'.format(file_name)
                lane1_tmp = []
                for j in range(binary_image1.shape[1]):
                    if binary_image1[i, j] == 50:
                        lane1_tmp.append(j)
                if len(lane1_tmp) > 0:
                    lane1_tmp = int(np.mean(np.array(lane1_tmp)))
                    lane1.append(lane1_tmp)
                else:
                    lane1.append(-2)

                lane2_tmp = []
                for j in range(binary_image2.shape[1]):
                    if binary_image2[i, j] == 100:
                        lane2_tmp.append(j)
                if len(lane2_tmp) > 0:
                    lane2_tmp = int(np.mean(np.array(lane2_tmp)))
                    lane2.append(lane2_tmp)
                else:
                    lane2.append(-2)

                lane3_tmp = []
                for j in range(binary_image3.shape[1]):
                    if binary_image3[i, j] == 150:
                        lane3_tmp.append(j)
                if len(lane3_tmp) > 0:
                    lane3_tmp = int(np.mean(np.array(lane3_tmp)))
                    lane3.append(lane3_tmp)
                else:
                    lane3.append(-2)

                lane4_tmp = []
                for j in range(binary_image4.shape[1]):
                    if binary_image4[i, j] == 200:
                        lane4_tmp.append(j)
                if len(lane4_tmp) > 0:
                    lane4_tmp = int(np.mean(np.array(lane4_tmp)))
                    lane4.append(lane4_tmp)
                else:
                    lane4.append(-2)
                h_samples.append(i)
                # cv2.line(binary_image, (0, i), (binary_image.shape[1], i), (255,), 1)

            lanes = [lane1, lane2, lane3, lane4]
            raw_file = base_name + '/' + file_name
            json_dict = {'h_samples': h_samples, 'lanes': lanes, 'raw_file': raw_file}
            json_line = json.dumps(json_dict)
            json_ff.write(json_line + '\n')
            print(img_fn)
            print(ann_fn)

            png_opath = os.path.join(frames_dir, os.path.splitext(file_name)[0] + '.png')
            # img4 = cv2.cvtColor(img4, cv2.COLOR_BGR2GRAY)
            cv2.imwrite(os.path.join(frames_dir, file_name), img4)
            cv2.imwrite(png_opath, binary_image_v2)
            # shutil.copy(ann_fn, os.path.join(gray_dir, file_name + '_Gray.lines.txt'))

            cv2.imshow('img', img)
            cv2.imshow('binary_image_v2', binary_image4)
            cv2.waitKey(10)
    json_ff.close()


def write_txt(fp):
    sub_dirs = os.listdir(fp)
    with open(os.path.join(fp, 'list', 'train.txt'), 'w') as f:
        for sub_dir in sub_dirs:
            img_names = os.listdir(os.path.join(fp, sub_dir))
            for img_name in img_names:
                if os.path.splitext(img_name)[1] == '.jpg':
                    to_out = '/' + sub_dir + '/' + img_name
                    f.write(to_out + '\r')


def to_vid(path):
    img_fns = os.listdir(path)
    four_cc = cv2.VideoWriter_fourcc(*"XVID")
    out = cv2.VideoWriter("output.avi", four_cc, 15.0, (1280, 720))

    for img_fn in img_fns:
        if os.path.splitext(img_fn)[1] == '.jpg':
            img = cv2.imread(os.path.join(path, img_fn))
            cv2.imshow('img', img)
            cv2.waitKey(10)
            out.write(img)

    out.release()


def bgr2gray(path):
    img_fns = os.listdir(path)
    for img_fn in img_fns:
        if os.path.splitext(img_fn)[1] == '.jpg':
            img = cv2.imread(os.path.join(path, img_fn))
            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
            cv2.imwrite(os.path.join(path, img_fn), img)
            cv2.imshow("img", img)
            cv2.waitKey(10)


def to720p(path):
    img_fns = os.listdir(path)
    for img_fn in img_fns:
        if os.path.splitext(img_fn)[1] == '.jpg':
            img = cv2.imread(os.path.join(path, img_fn))
            img = cv2.resize(img, (1280, 720))
            cv2.imwrite(os.path.join(path, img_fn), img)
            cv2.imshow("img", img)
            cv2.waitKey(10)


if __name__ == '__main__':
    # to720p(r'C:\dataset\A2RL\yas_lane_tusimple_240417\camfr2')
    # to_vid(r'C:\dataset\A2RL\yas_lane_tusimple_240414_2\Yas_Grey_240406')
    # show_lane_in_dir(r'C:\dataset\A2RL\yas_lane_tusimple_240417\Yas_Gray_240419')
    # write_txt('C:/dataset/A2RL/ws_lane_wg_240403')
    bgr2gray(r'C:\dataset\A2RL\BB-RaceCar-Det-v2404\TeamCar-0414\scaled_images')
